#!/usr/bin/env python

import os
import re
import sys
import gzip
import cStringIO

from create_gen_desc_db import parse_cds, half_md5


def parse_genes(s):
    f = cStringIO.StringIO(s)
    gene_started = False
    borders, ncbi_id = None, None
    for line in f:
        if gene_started:
            if not line.startswith('                     '):
                if borders and ncbi_id:
                    yield borders, ncbi_id
                    borders, ncbi_id = None, None
                    gene_started = False
                else:
                    raise Exception('Bad CDS')
            else:
                line = line.strip()
                if line.startswith("/db_xref=\"GeneID:"):
                    ncbi_id = line.split("\"")[1].split(":")[1]
        elif line.startswith('     gene            '):
            borders = line.strip().split()[1]
            gene_started = True


if __name__ == "__main__":
 
    orgs_dir        = sys.argv[1]
    gene_info_fname = sys.argv[2]

    cnt_all = 0
    cnt_genes = 0
    gene_info_file = open(gene_info_fname, 'w')

    for locus in os.listdir(orgs_dir):
        fobj = gzip.open(os.path.join(orgs_dir, locus))
        raw_desc = fobj.read()
        fobj.close()
        cnt_all += 1
        cds_borders = dict(parse_cds(raw_desc))
        for borders, ncbi_id in parse_genes(raw_desc):
            gene_id = "%s_%s" % (locus, half_md5(borders))
            product = cds_borders.get(borders, '__NULL__').strip()
            gene_info_file.write("%s\t%s\t%s\t%s\n" % (gene_id, ncbi_id, borders, product))
            cnt_genes += 1

    gene_info_file.close()

    print "Done! (%s/%s)" % (cnt_all, cnt_genes)

